// AppInitDll.cpp : Defines the entry point for the DLL application.
//

#include "stdafx.h"
#include "appinitdll.h"

#ifdef _MANAGED
#pragma managed(push, off)
#endif

//
// Registry locations
//

#define szAPP_INIT_DLLS_KEY		"Software\\Microsoft\\Windows NT\\CurrentVersion\\Windows"
#define szAPP_INIT_DLLS_VALUE	"AppInit_Dlls"
#define szDLL_NAME				"appinitdll.dll"

#define szIMAGE_IDATA_SECTION	".idata"
#define szWINSCARD_DLL          "winscard.dll"
#define szSCARD_TRANSMIT        "SCardTransmit"

// Convert a relative virtual address to a virtual address
#define RVA_TO_VAD(_module, _rva) ((DWORD) _module + _rva)

/*

//
// Print out hex data to console
//
#define CROW                                                16
static void PrintBytes(LPSTR szTitle, PBYTE pbData, DWORD cbData)
{
    DWORD iByte = 0, iRowItem = 0;
    CHAR rgbRow [100];

	StringCbPrintfA(
        rgbRow, sizeof(rgbRow), "\n%s -- %d bytes\n", szTitle, cbData);
    OutputDebugStringA(rgbRow);

    while (iByte < cbData)
    {
        for (   iRowItem = 0; 
                iRowItem < CROW && iByte < cbData; 
                iRowItem++, iByte++)
            StringCbPrintfA(
                rgbRow + (iRowItem * 3), 
                sizeof(rgbRow) - (iRowItem * 3),
                "%02X ", 
                pbData[iByte]);

        OutputDebugStringA(rgbRow);
        OutputDebugStringA("\n");
    }
}

//
// Test replacement routine
//
LONG 
WINAPI
_MySCardTransmit(
    IN SCARDHANDLE hCard,
    IN LPCSCARD_IO_REQUEST pioSendPci,
    IN LPCBYTE pbSendBuffer,
    IN DWORD cbSendLength,
    IN OUT LPSCARD_IO_REQUEST pioRecvPci,
    OUT LPBYTE pbRecvBuffer,
    IN OUT LPDWORD pcbRecvLength)
{
    LONG lReturn = SCardTransmit(
        hCard, 
        pioSendPci,
        pbSendBuffer,
        cbSendLength,
        pioRecvPci,
        pbRecvBuffer,
        pcbRecvLength);

    PrintBytes("Transmit / Send", (PBYTE) pbSendBuffer, cbSendLength);
    PrintBytes("Transmit / Receive", pbRecvBuffer, *pcbRecvLength);

    return lReturn;
}
*/

//
// Patch the IAT of an image in memory
// 
// If the target image is the host exe, set szImageToPatch = NULL and 
// fImageIsAnExe = TRUE.
//
DWORD
WINAPI
PatchImage(
    __in_opt        LPSTR szImageToPatch,
    __in            BOOL fImageIsAnExe,
	__in	        LPSTR szImportDll,
	__in	        LPSTR szImportRoutine,
	__in	        PROC pfnNewRoutine)
{
	DWORD status = 0;
	HMODULE hImageToPatch = 0;
	HMODULE hTargetDll = 0;
	FARPROC pfnTargetRoutine = NULL;
	PROC *ppfnTemp = NULL;
	PIMAGE_DOS_HEADER pDosHeader = NULL;
	PIMAGE_NT_HEADERS pNtHeader = NULL;
	PIMAGE_SECTION_HEADER rgSectionHeaders = NULL;
	PIMAGE_IMPORT_DESCRIPTOR rgImportDescriptors = NULL;
	PIMAGE_THUNK_DATA rgThunkDataHints = NULL;
	PIMAGE_THUNK_DATA rgThunkDataFixups = NULL;
	DWORD iSection = 0;
	DWORD iImport = 0;
	DWORD iThunk = 0;
	MEMORY_BASIC_INFORMATION mbi = {0};

	__try
	{
		//
		// Get a function ptr for the target export
		//

		if (NULL == (hTargetDll = GetModuleHandleA(szImportDll)))
		{
			status = GetLastError();
			__leave;
		}

		if (NULL == (pfnTargetRoutine = GetProcAddress(
			hTargetDll, szImportRoutine)))
		{
			status = GetLastError();
			__leave;
		}

		//
		// Get a module handle for the image to patch
		//

		if (NULL == (hImageToPatch = GetModuleHandleA(szImageToPatch)))
		{
			status = GetLastError();
			__leave;
		}

		//
		// Module handle value is the DOS header ptr
		//

		pDosHeader = (PIMAGE_DOS_HEADER) hImageToPatch;

		//
		// Find the NT headers ptr
		//

		pNtHeader = (PIMAGE_NT_HEADERS) RVA_TO_VAD(
			hImageToPatch, pDosHeader->e_lfanew);

        if (TRUE == fImageIsAnExe)
        {
            //
            // Find the array of section headers
            //

            rgSectionHeaders = IMAGE_FIRST_SECTION(pNtHeader);

            //
            // Find the .idata section
            //

            for (	iSection = 0; 
                    iSection < pNtHeader->FileHeader.NumberOfSections; 
                    iSection++)
                if (0 == _strnicmp(
                        (const char *) rgSectionHeaders [iSection].Name, 
                        szIMAGE_IDATA_SECTION,
                        IMAGE_SIZEOF_SHORT_NAME))
                    break;

            // Verify we found the section
            if (iSection == pNtHeader->FileHeader.NumberOfSections)
            {
                status = ERROR_NOT_FOUND;
                __leave;
            }

            //
            // Find the beginning of the image import descriptor array
            //

            rgImportDescriptors = (PIMAGE_IMPORT_DESCRIPTOR) RVA_TO_VAD(
                hImageToPatch, rgSectionHeaders [iSection].VirtualAddress);
        }
        else
        {
            rgImportDescriptors = (PIMAGE_IMPORT_DESCRIPTOR) RVA_TO_VAD(
                hImageToPatch, 
                pNtHeader->OptionalHeader.DataDirectory [IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
        }

		//
		// Find the entry for the requested dll
		//

		for (iImport = 0; 0 != rgImportDescriptors [iImport].Name; iImport++)
			if (0 == _stricmp(
					(LPSTR) RVA_TO_VAD(
						hImageToPatch, rgImportDescriptors [iImport].Name),
					szImportDll))
				break;

		// Verify we found the import entry
		if (0 == rgImportDescriptors [iImport].Name)
		{
			status = ERROR_NOT_FOUND;
			__leave;
		}

		//
		// Find the entry for the requested Dll export
		//

		rgThunkDataHints = (PIMAGE_THUNK_DATA) RVA_TO_VAD(
			hImageToPatch, rgImportDescriptors [iImport].OriginalFirstThunk);
		rgThunkDataFixups = (PIMAGE_THUNK_DATA) RVA_TO_VAD(
			hImageToPatch, rgImportDescriptors [iImport].FirstThunk);

		for (	iThunk = 0; 
				NULL != rgThunkDataFixups [iThunk].u1.Function; 
				iThunk++)
			if (	pfnTargetRoutine == (FARPROC) 
					rgThunkDataFixups [iThunk].u1.Function)
				break;

		// Verify we found the export
		if (NULL == rgThunkDataFixups [iThunk].u1.Function)
		{
			status = ERROR_NOT_FOUND;
			__leave;
		}

		ppfnTemp = (PROC *) &(rgThunkDataFixups [iThunk].u1.Function);

		//
		// Mark the VM read-write
		//

		if (0 == VirtualQuery(ppfnTemp, &mbi, sizeof(mbi)))
		{
			status = ERROR_INTERNAL_ERROR;
			__leave;
		}

		if (FALSE == VirtualProtect(
			mbi.BaseAddress, mbi.RegionSize, PAGE_READWRITE, &mbi.Protect))
		{
			status = GetLastError();
			__leave;
		}

		//
		// Patch the import with our new "export"
		//

		*ppfnTemp = pfnNewRoutine;

		//
		// Mark the VM back to its original protection
		//

		if (FALSE == VirtualProtect(
			mbi.BaseAddress, mbi.RegionSize, mbi.Protect, &mbi.Protect))
		{
			status = GetLastError();
			__leave;
		}
	}
	__finally
	{
		// Nothing to do
	}

	return status;
}

//
// Loader entry point for this Dll
//
BOOL 
APIENTRY 
DllMain( 
	__in	HMODULE hModule,
	__in    DWORD ul_reason_for_call,
	__in    LPVOID lpReserved)
{
	switch (ul_reason_for_call)
	{
	case DLL_PROCESS_ATTACH:

		//
		// Don't need to receive Thread Attach/Detach.
		//

		DisableThreadLibraryCalls(hModule);

		//
		// Test patching 
		//

        /*
		PatchImage(
            NULL,
            szWINSCARD_DLL, 
            szSCARD_TRANSMIT, 
            (PROC) _MySCardTransmit);
        */
		break;

	case DLL_PROCESS_DETACH:

		//
		// Nothing to do 
		//

		break;
	}

    return TRUE;
}

//
// Helper routine to open the app init dlls reg key
//
DWORD
WINAPI
_OpenAppInitDllsRegKey(
	__out	PHKEY phKey)
{
	return RegOpenKeyExA(
		HKEY_LOCAL_MACHINE,
		szAPP_INIT_DLLS_KEY,
		0,
		KEY_SET_VALUE,
		phKey);
}

//
// Add registry info
//
void
APIENTRY
DllRegisterServer(void)
{
	DWORD status = ERROR_SUCCESS;
	HKEY hKey = 0;

	__try
	{
		//
		// Open the app dlls reg key
		//

		status = _OpenAppInitDllsRegKey(&hKey);
		if (ERROR_SUCCESS != status)
			__leave;

		//
		// Set this dll as the value
		//

		status = RegSetValueExA(
			hKey,
			szAPP_INIT_DLLS_VALUE,
			0,
			REG_SZ,
			(const BYTE *) szDLL_NAME,
			(DWORD) (sizeof(CHAR) * (strlen(szDLL_NAME) + 1)));
		if (ERROR_SUCCESS != status)
			__leave;
	}
	__finally
	{
		if (0 != hKey)
			RegCloseKey(hKey);
	}
}

//
// Remove registry info
//
void
APIENTRY
DllUnregisterServer(void)
{
	DWORD status = ERROR_SUCCESS;
	HKEY hKey = 0;

	__try
	{
		//
		// Open the app dlls reg key
		//

		status = _OpenAppInitDllsRegKey(&hKey);
		if (ERROR_SUCCESS != status)
			__leave;

		//
		// Clear the value
		//

		status = RegSetValueExA(
			hKey,
			szAPP_INIT_DLLS_VALUE,
			0,
			REG_SZ,
			(const BYTE *) "",
			sizeof(CHAR));
		if (ERROR_SUCCESS != status)
			__leave;
	}
	__finally
	{
		if (0 != hKey)
			RegCloseKey(hKey);
	}
}

#ifdef _MANAGED
#pragma managed(pop)
#endif

